"""
向 llm_models 表中插入测试数据
"""

import uuid
from sqlalchemy.orm import Session
from models import LLMModel, ModelModality, ModelStatus, User
from dao.postgresql_connect import SessionLocal


def seed_llm_models():
    """向 llm_models 表中插入测试数据"""
    db = SessionLocal()
    try:
        # 检查是否已有数据
        existing_models = db.query(LLMModel).count()
        if existing_models > 0:
            print(f"数据库中已有 {existing_models} 条模型数据，跳过插入测试数据")
            return

        # 获取管理员用户
        admin_user = db.query(User).filter(User.role == "ADMIN").first()
        if not admin_user:
            print("未找到管理员用户，请先创建管理员用户")
            return

        admin_id = admin_user.id
        print(f"使用管理员用户 ID: {admin_id} 作为创建者和更新者")

        # 创建测试数据
        test_models = [
            {
                "name": "DeepSeek V3 Base (free)",
                "model_id": "deepseek/deepseek-v3-base:free",
                "provider": "OpenRouter",
                "modality": ModelModality.TEXT_ONLY,
                "base_url": "https://openrouter.ai/api/v1",
                "api_key": "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",
                "status": ModelStatus.ACTIVE,
                "description": "深度求索的开源大模型，免费使用，支持流式输出，多人使用会卡顿",
                "config": {
                    "max_tokens": 4096,
                    "temperature": 0.7
                }
            },
            {
                "name": "DeepSeek V3 0324 (free)",
                "model_id": "deepseek/deepseek-chat-v3-0324:free",
                "provider": "OpenRouter",
                "modality": ModelModality.TEXT_ONLY,
                "base_url": "https://openrouter.ai/api/v1",
                "api_key": "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",
                "status": ModelStatus.ACTIVE,
                "description": "深度求索的开源大模型，较前版本编程能力有较大提升，免费使用，支持流式输出，多人使用会卡顿",
                "config": {
                    "max_tokens": 4096,
                    "temperature": 0.7
                }
            },
            {
                "name": "DeepSeek: DeepSeek R1 Zero (free)",
                "model_id": "deepseek/deepseek-r1-zero:free",
                "provider": "OpenRouter",
                "modality": ModelModality.TEXT_ONLY,
                "base_url": "https://openrouter.ai/api/v1",
                "api_key": "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",
                "status": ModelStatus.ACTIVE,
                "description": "深度求索的开源推理模型，目前较好的推理模型，免费使用。",
                "config": {
                    "max_tokens": 128000,
                    "temperature": 0.7
                }
            },
            {
                "name": "Quasar Alpha",
                "model_id": "openrouter/quasar-alpha",
                "provider": "OpenRouter",
                "modality": ModelModality.TEXT_IMAGE,
                "base_url": "https://openrouter.ai/api/v1",
                "api_key": "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",
                "status": ModelStatus.ACTIVE,
                "description": "OpenRouter 的 Quasar Alpha 模型，支持文本和图像生成，免费使用，支持流式输出，多人使用会卡顿",
                "config": {
                    "max_tokens": 128000,
                    "temperature": 0.7
                }
            },
            {
                "name": "Qwen2.5 VL 32B Instruct (free)",
                "model_id": "qwen/qwen2.5-vl-32b-instruct:free",
                "provider": "OpenRouter",
                "modality": ModelModality.MULTIMODAL,
                "base_url": "https://openrouter.ai/api/v1",
                "api_key": "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",
                "status": ModelStatus.ACTIVE,
                "description": "通义千问的开源多模态模型，支持图像、文本、视频输入，免费使用。",
                "config": {
                    "max_tokens": 8000,
                    "temperature": 0.7
                }
            }
        ]

        # 插入测试数据
        for model_data in test_models:
            # 创建模型实例
            model = LLMModel(
                **model_data,
                created_by=admin_id,
                updated_by=admin_id
            )

            # 添加到数据库
            db.add(model)

        # 提交事务
        db.commit()
        print(f"成功插入 {len(test_models)} 条测试数据到 llm_models 表")

    except Exception as e:
        db.rollback()
        print(f"插入测试数据失败: {str(e)}")
    finally:
        db.close()


if __name__ == "__main__":
    seed_llm_models()
